﻿#define _CRT_SECURE_NO_WARNINGS 1
#include<iostream>
#include<assert.h>
#include<vector>
#include<utility>
using namespace std;


// 枚举值表示颜色
enum Colour
{
	RED,
	BLACK
};


// 这里我们默认按key/value结构实现
template<class K, class V>
struct RBTreeNode
{
	// 这里更新控制平衡也要加入parent指针
	pair<K, V> _kv;
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	Colour _col;


	RBTreeNode(const pair<K, V>& kv)
		:_kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
	{}
};


template<class K, class V>
class RBTree
{
	typedef RBTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv) {
		//前面和搜索二叉树的的规则一样插入数值即可
		Node* cur = _root;
		Node* parent = _root;

		while (cur) {
			if (kv.first > cur->_kv.first) {
				parent = cur;
				cur = cur->_right;
			}
			else if(kv.first < cur->_kv.first){
				parent = cur;
				cur = cur->_left;
			}
			else {
				//cout << "不能插入相同的值" << endl;
				return false;
			}
		}
		Node* newroot = new Node(kv);
		newroot->_col = RED;
		if (_root == nullptr) {
			_root = newroot;
			_root->_col = BLACK;
			return true;
		}
		else {
			if (kv.first > parent->_kv.first) {
				parent->_right = newroot;
				newroot->_parent = parent;
			}
			else {
				parent->_left = newroot;
				newroot->_parent = parent;
			}
		}
		//开始平衡红黑树的规则
		cur = newroot;Node* uncle;Node* grandparent = parent->_parent;
		while (cur->_col == RED&&cur->_parent && cur->_parent->_col == RED) {
			parent = cur->_parent;
			if (parent->_parent->_kv.first > parent->_kv.first) {
				uncle = parent->_parent->_right;
			}
			else {
				uncle = parent->_parent->_left;
			}
			grandparent = parent->_parent;
			//情况一 叔叔存在且为红
			if (uncle && uncle->_col == RED) {
				uncle->_col = BLACK;
				parent->_col = BLACK;
				grandparent->_col = RED;
			}
			else {
				//情况2 叔叔不在
				if (!uncle) {
					//说明此为新加入的节点cur
					//判断单旋还是双旋转
					if (grandparent->_kv.first > parent->_kv.first) {
						if (parent->_kv.first > cur->_kv.first) {
							//单右旋
							RotateR(grandparent);

							parent->_col = BLACK;
							grandparent->_col = RED;
						}
						else {
							//左右双旋
							RotateLR(grandparent);
							cur->_col = BLACK;
							grandparent->_col = RED;
							parent->_col = RED;
						}
					}
					else {
						if (parent->_kv.first > cur->_kv.first) {
							//右左双旋
							RotateRL(grandparent);
							cur->_col = BLACK;
							grandparent->_col = RED;
							parent->_col = RED;
						}
						else {
							//单左旋
							RotateL(grandparent);
							parent->_col = BLACK;
							grandparent->_col = RED;
						}
					}
				}else {
					//情况 3叔叔存在且为黑
					//说明这不是新加入的节点，是旧节点改过来的
					//判断单旋还是双旋转
					if (grandparent->_kv.first > parent->_kv.first) {
						if (parent->_kv.first > cur->_kv.first) {
							//单右旋
							RotateR(grandparent);
							parent->_col = BLACK;
							grandparent->_col = RED;
						}
						else {
							//左右双旋
							RotateLR(grandparent);
							cur->_col = BLACK;
							grandparent->_col = RED;
							uncle->_col = BLACK;
							parent->_col = RED;
						}
					}
					else {
						if (parent->_kv.first > cur->_kv.first) {
							//右左双旋
							RotateRL(grandparent);
							cur->_col = BLACK;
							grandparent->_col = RED;
							uncle->_col = BLACK;
							parent->_col = RED;
						}
						else {
							//单左旋
							RotateL(grandparent);
							parent->_col = BLACK;
							grandparent->_col = RED;
						}
					}
				}
			}
			

			cur = grandparent;
		}

		_root->_col = BLACK;
		return true;
	}

	// 右单旋
	void RotateR(Node* pParent) {
		Node* subl = pParent->_left;
		Node* sublr = subl->_right;
		subl->_right = pParent;
		if (pParent->_parent) {//判断根节点
			subl->_parent = pParent->_parent;
			if (pParent->_parent->_kv.first > pParent->_kv.first) {
				pParent->_parent->_left = subl;
			}
			else pParent->_parent->_right = subl;
		}
		else {
			subl->_parent = nullptr;
			_root = subl;
		}

		if (sublr) sublr->_parent = pParent;
		pParent->_parent = subl;
		subl->_right = pParent;
		pParent->_left = sublr;
	}
	// 左单旋
	void RotateL(Node* pParent) {
		Node* subr = pParent->_right;
		Node* subrl = subr->_left;
		subr->_left = pParent;
		if (pParent->_parent) {//判断根节点
			subr->_parent = pParent->_parent;
			if (pParent->_parent->_kv.first > pParent->_kv.first) {
				pParent->_parent->_left = subr;
			}
			else pParent->_parent->_right = subr;
		}
		else {
			subr->_parent = nullptr;
			_root = subr;
		}

		if (subrl) {
			subrl->_parent = pParent;
		}
		pParent->_parent = subr;
		subr->_left = pParent;
		pParent->_right = subrl;
	}
	// 右左双旋
	void RotateRL(Node* pParent) {
		Node* subr = pParent->_right;
		//Node* subrl = subr->_left;
		RotateR(subr);
		RotateL(pParent);
	}
	// 左右双旋
	void RotateLR(Node* pParent) {
		Node* subl = pParent->_left;
		//Node* sublr = pParent->_left->_right;
		RotateL(subl);
		RotateR(pParent);
	}

	void InOrder() {
		_InOrder(_root);
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr)
		{
			return;
		}


		_InOrder(root->_left);
		cout << root->_kv.first << " ";
		_InOrder(root->_right);
	}


	// 前序递归遍历
	bool Check(Node* root, int blackNum, const int refNum)
	{
		if (root == nullptr)
		{
			// 前序遍历走到空时，意味着一条路径走完了
			//cout << blackNum << endl;
			if (refNum != blackNum)
			{
				cout << "存在黑色结点的数量不相等的路径" << endl;
				return false;
			}
			return true;
		}


		// 检查孩子不太方便，因为孩子有两个，且不一定存在，反过来检查父亲就方便多了
		if (root->_col == RED && root->_parent && root->_parent->_col == RED)
		{
			cout << root->_kv.first << "存在连续的红色结点" << endl;
			return false;
		}


		if (root->_col == BLACK)
		{
			blackNum++;
		}


		return Check(root->_left, blackNum, refNum)
			&& Check(root->_right, blackNum, refNum);
	}

	Node* Find(const K& key)
	{
		Node* cur = _root;
			while (cur)
			{
				if (cur->_kv.first < key)
				{
					cur = cur->_right;
				}
				else if (cur->_kv.first > key)
				{
					cur = cur->_left;
				}
				else
				{
					return cur;
				}
			}
		return nullptr;
	}

	bool IsBalanceTree()
	{
		if (_root == nullptr)
			return true;


		if (_root->_col == RED)
			return false;


		// 参考值
		int refNum = 0;
		Node* cur = _root;
		while (cur)
		{
			if (cur->_col == BLACK)
			{
				++refNum;
			}
			cur = cur->_left;
		}


		return Check(_root, 0, refNum);
	}


private:
	Node* _root = nullptr;
};




// 测试代码
void TestRBTree1()
{
	RBTree<int,int> t;
	// 常规的测试⽤例
	//int a[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
	// 特殊的带有双旋场景的测试⽤例
	int a[] = { 4, 2,5,5,5,5,5,5,5, 6, 1, 3, 5, 15, 7, 16, 14 };
	for (auto e : a)
	{
		t.Insert(make_pair(e,e));
	}
	t.InOrder();
	if (t.IsBalanceTree()) {
		cout << "TRUE" << endl;
	}
	else
	{
		cout << "FALSE" << endl;

	}

}
// 插⼊⼀堆随机值，测试平衡，顺便测试⼀下⾼度和性能等
void TestRBTree2()
{
	const int N = 100000;
	vector<int> v;
	v.reserve(N);
	srand(time(0));
	for (size_t i = 0; i < N; i++)
	{
		v.push_back(rand() + i);
	}
	size_t begin2 = clock();
	RBTree<int,int> t;
	for (auto e : v)
	{
		t.Insert(make_pair(e,e));
	}
	size_t end2 = clock();
	cout << "Insert:" << end2 - begin2 << endl;
	if (t.IsBalanceTree()) {
		cout << "TRUE" << endl;
	}
	else
	{
		cout << "FALSE" << endl;

	}
	//cout << "Height:" << t.Height() << endl;
	//cout << "Size:" << t.Size() << endl;
	//`size_t begin1 = clock();
	// 确定在的值
	for (auto e : v)
	{
	t.Find(e)->_kv.second;

	}
	// 随机值
	/*for (size_t i = 0; i < N; i++)
	{
		t.Find((rand() + i));
	}*/
	/*size_t end1 = clock();
	cout << "Find:" << end1 - begin1 << endl;*/
}

int main() {
	TestRBTree2();
	TestRBTree1();
}